import os
import sys
import cv2
import dlib
import shutil
import corecode.api as corecode
import numpy as np
np.set_printoptions(threshold=32, edgeitems=10)
try:
   import cPickle as pickle
except:
   import pickle

def getAllArgs(howOut="return"):
    allArgs = sys.argv
    allArgs = allArgs[1:]
    return allArgs

def folderContents(dire):
    for dirpath, dirnames, filenames in os.walk(dire):
        return dirpath, dirnames, filenames

def identFacesInPictures(names, known_faces, main_image):
    doneList = []               
    face_locations = corecode.face_locations(main_image)
    face_encodings = corecode.face_encodings(main_image, face_locations)    
    for i, ((top, right, bottom, left), atv_encoding) in enumerate(zip(face_locations, face_encodings)):
        for ii, (name, known_face) in enumerate(zip(names, known_faces)): 
           if (str(top) + str(left)) in doneList: break;

           match = corecode.compare_faces([known_face], atv_encoding)
              
           if match[0] == True:
                cv2.rectangle(main_image, (left, top), (right, bottom), (255, 0, 0), 2) 
                # Draw a label with a name below the face
                cv2.rectangle(main_image, (left, bottom - 35), (right, bottom), (255, 0, 0), cv2.FILLED)
                font = cv2.FONT_HERSHEY_DUPLEX
                cv2.putText(main_image, name, (left + 6, bottom - 6), font, 1.0, (0, 0, 0), 1)

                doneList.append(str(top) + str(left))

           elif match[0] == False:
                name = "Other"
                cv2.rectangle(main_image, (left, top), (right, bottom), (0, 0, 255), 2) 
                # Draw a label with a name below the face
                cv2.rectangle(main_image, (left, bottom - 35), (right, bottom), (0, 0, 255), cv2.FILLED)
                font = cv2.FONT_HERSHEY_DUPLEX
                cv2.putText(main_image, name, (left + 6, bottom - 6), font, 1.0, (255, 255, 255), 1)                    
    return main_image


def loadPickle(ff):
    with open(ff, "rb") as openfile:
         pData = (pickle.load(openfile, encoding='latin1'))  
         return pData
    ff.close()

def main():
    mainNameList = []
    mainEmbeList = []

    cwd = os.getcwd()

    print ("Loading embeddings....")
    mainLoc, fConts, fFiles = folderContents(cwd + "/embeddings/")
    for i in fConts: 
        mainLoc_2, fConts_2, fFiles_2 = folderContents(mainLoc + "/" + i)
        fFiles_2b = filter(lambda x : x != '*.jpg', fFiles_2)
        for i in fFiles_2b:
            if not any(value in (i.lower()) for value in (".jpeg", ".jpg", ".gif", ".png")):
               totalString = (mainLoc_2 + "/" + i)

               name = totalString.split("//")[-1].split("/")[0]
               nameFL = name.replace(".", " ").title()
               eT = totalString.split("//")[-1].split(".")[1].split("/")[1]
               nameMain = (nameFL + " - EmbeddingID[{}]".format(eT))
               mainNameList.append(nameMain)

               ebData = loadPickle(totalString)
               mainEmbeList.append(ebData)

               print  (nameMain + str(ebData) + '\n')
    print ("Done loading embeddings.")

    gAA = getAllArgs()
    mA = gAA[0]
    tMain_image = corecode.load_image_file(mA)
    tFI = identFacesInPictures(mainNameList, mainEmbeList, tMain_image)
    cv2.namedWindow(str(mA), cv2.WINDOW_NORMAL)
    cv2.imshow(str(mA),  cv2.cvtColor(tFI, cv2.COLOR_BGR2RGB))
    cv2.resizeWindow(str(mA), 1000, 1000)
    cv2.waitKey(0)


main()









